from packetholder import packetholder
from packet import packet
from agent import agent
from constants import SYSTEM, PRIO, RCONSTS

LINK_N = 0
LINK_S = 1
LINK_E = 2
LINK_W = 3

class router(agent):
  # Info for routing
  loc = (-1, -1)
  rtable = None
  
  # info on outputs
  link_bw_f = None
  outputs = None

  # router stats
  stat_all_util_bw = 0 # ratio, how much of router-router BW used
  stat_max_util_bw = 0 # ratio, max of all links for router-router BW used

  u_stat_all_util_bw = 0
  u_stat_max_util_bw = 0
  
  # buffers to store partial messages
  buffers = None

  num_in_buffers = 4 # recall, each router always gets 4 buffers for
  in_buffer_map = None

  def __init__(self, in_rtable):
    super(router, self).__init__()

    self.link_bw_f = [None] * 4
    self.outputs = [None] * 4

    self.buffers = [None] * PRIO.NUM_LEVELS
    for vn in range(PRIO.NUM_LEVELS):
      routes = [None] * 4
      for r in range(4):
        routes[r] = packetholder(RCONSTS.MAX_PACKETS, RCONSTS.MAX_FLITS)
      self.buffers[vn] = routes
    self.in_buffer_map = {} # do not bother with mapping routers

    self.rtable = in_rtable

  def add_out(self, agent_out, bw_f, direction=None):
    if(direction is None):
      self.outputs.append(agent_out)
      self.link_bw_f.append(bw_f)
    else:
      self.outputs[direction] = agent_out
      self.link_bw_f[direction] = bw_f

  def determine_route(self, pkt):
    # dimension order routing, x first
    d_coord = self.rtable.lookup(pkt.dest)
    if(d_coord[0] != self.loc[0]):
      if(d_coord[0] > self.loc[0]):
        return LINK_E
      else:
        return LINK_W
    elif(d_coord[1] != self.loc[1]):
      if(d_coord[1] > self.loc[1]):
        return LINK_N
      else:
        return LINK_S
    return None

  def accept(self, src_node, inc_pkt, max_trans_f, route_hint=None):
    # this system effectively uses input buffering for the routers
    #   so create a buffer for (almost) each input buffer
    routes = self.buffers[inc_pkt.priority]
    if(route_hint is None):
      if src_node.node_id not in self.in_buffer_map:
        for vn in self.buffers:
          new_ph = packetholder(RCONSTS.MAX_PACKETS, RCONSTS.MAX_FLITS)
          vn.append(new_ph)
        self.in_buffer_map[src_node.node_id] = self.num_in_buffers
        self.num_in_buffers += 1

      route = self.in_buffer_map[src_node.node_id]
    else:
      # 'route hints' come from other routers
      route = route_hint
    return routes[route].accept(inc_pkt, max_trans_f)

  def process(self):
    # prepare some link BW
    avail_f = list(self.link_bw_f)

    # prioritize higher priority packets
    for vn in reversed(self.buffers):
      # store all potential packets in output queues
      #   so we can sort and apply age-based local flow ctrl
      output_queues = [None] * len(self.outputs)
      for routes in vn:
        for pkt in routes.contents:
          # Make sure actually have part of packet to send
          if pkt.last_f > 0:
            # get the desird routing direction
            o_entry = self.determine_route(pkt)
            # check to see if already at destination router
            if o_entry is None:
              # find output from this router to desired node
              for i in range(len(self.outputs)):
                if self.outputs[i] and (self.outputs[i].node_id == pkt.dest):
                  o_entry = i
                  break
            if output_queues[o_entry] is None: output_queues[o_entry] = [] 
            output_queues[o_entry].append(pkt)
      
      # all packets for this level routed, now decide who gets to go
      for o_entry in range(len(self.outputs)):
        if(output_queues[o_entry] is not None):
          sorted_queue = sorted(output_queues[o_entry], key=packet.get_age)
          for pkt in sorted_queue:
            if(avail_f[o_entry] <= 0): break # Do not bother trying to transmit packets if no link BW available
            trans_f = self.transmit(self.outputs[o_entry], pkt, avail_f[o_entry], o_entry)
            avail_f[o_entry] -= trans_f
    
    # calculate statistics
    total_bw = 0
    used_bw = 0
    for i in range(4):
      if avail_f[i] is not None:
        link_used = self.link_bw_f[i] - avail_f[i]
        used_bw += link_used
        total_bw += self.link_bw_f[i]
    self.u_stat_all_util_bw = used_bw / (total_bw * 1.0)
    self.u_stat_max_util_bw = max(self.stat_max_util_bw, self.u_stat_all_util_bw)


  def commit(self):
    for vn in self.buffers:
      for route in vn:
        route.commit()
    self.stat_all_util_bw = self.u_stat_all_util_bw
    self.stat_max_util_bw = self.u_stat_max_util_bw

  # this function is for debugging purposes
  def to_string(self):
    return "%0.2f" % self.stat_all_util_bw
    accum = 0
    for vn in self.buffers:
      for route in vn:
#        accum = max(accum,len(route.contents))
        for pkt in route.contents:
          accum += pkt.last_f - pkt.first_f
    return str(accum)
